"""
#===============================================================================
# License
#===============================================================================
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

#===============================================================================
Copyright (c) 2009 bryan p. conrad
bconrad2@gmail.com

Repository Path:    
$$HeadURL$$

Created on:           7/27/2009
Last Revised on:      $$Date$$
Version:              $$Revision$$
ID:                   $$Id$$
Author:               $$Author$$ :: (bryan)

#===============================================================================


#===============================================================================
# Documentation
#===============================================================================

"""

#===============================================================================
# Imports
#===============================================================================
import JT
import pylab


#===============================================================================
# Generic Data Structure
#===============================================================================
class DataNode(list):
    """
    An abstract definition for each node on the Data tree structure.
    """
#    def __init__(self):
#        self._attrs = {}    # attributes and child elements
#        self.data = None    # child text data
#    def __len__(self):
#        # treat single element as a list of 1
#        return 1
#    def __getitem__(self, key):
#        if isinstance(key, basestring):
#            return self._attrs.get(key,None)
#        else:
#            return [self][key]
#    def __contains__(self, name):
#        return self._attrs.has_key(name)
#    def __nonzero__(self):
#        return bool(self._attrs or self.data)
#    def __getattr__(self, name):
#            if name.startswith('__'):
#                # need to do this for Python special methods???
#                raise AttributeError(name)
#            return self._attrs.get(name,None)
#
#    def add(self, object):
#        """ Add an object to the list at this node.
#        Returns the  index of the value in the list, if only one value exists 
#        returns 0. 
#        """
##        if name in self._attrs:
#            # multiple attribute of the same name are represented by a list
##            children = self._attrs[name]
##            if not isinstance(children, list):
##                children = [children]
##                self._attrs[name] = children
##            children.append(value)
##            self.__setattr__(name,children)
#        try:
#            self.append(object)
#            return len(self) - 1    #subtract one to get 0 based index
#        except Exception:
#            print "Error setting :", object
#            raise
#
#        
#    def __str__(self):
#        return self.data or ''
    def __str__(self):
        #str='Nodes:'
        tree = 'ROOT\n'
        for k in self.__dict__:
            tree = tree + '|__%s' % (str(k))
            try:
                tree = tree + self.get_progeny(self.__getattribute__(k))
            except:
                tree = tree + '\n   |__  %s' % (str(k))
                
            tree = tree + '\n'
        return tree
            #str = str + '\n    |__' + self.get_progeny(self.__getattribute__(k))
    
    
    def get_progeny(self, parent): 
        out = ''
        for k in parent.__dict__:
            out = out + '\n   |__  ' + str(k)
        return out
    
    def len(self):
        """ Returns the length of the current list
        """
        return len(self)
    
    def add(self):
        """ appends a new instance of the base class to the internal list.
        The object can be retrieved using the signature DataNode[y]
        """
        self.append(self.__class__())
        
#    def __getitem__(self,y):
#        if type(super(DataNode,self).__getitem__(y)) is list:
#            return super(DataNode,self).__getitem__(y)[0]
#        else:
#            return super(DataNode,self).__getitem__(y)
        



class Data(object):
    """
    A container for JointTrack data.  Data is initialized with one frame.
    """
    def __init__(self):
        self.Frames = Frames()
        #self.Frames.add()
        # Initialize with one Frame (minimum required to do something useful)
        #self.Frames.append(Frames())
        self.currentFrame = 0
        
    
class Frames(DataNode):
    """
    A Frame is a slice in time, it contains:
        volumes (potentially several)
        fixed images[s] (one for each fluoroscopic view)
        poses: one for each model
        display settings: parameter for setting the display of each volume/image        
    """
    def __init__(self):
        """
        """
        super(DataNode,self).__init__()
        self.Volumes = Volumes()
        self.Volumes.add()
        self.FixedImages = FixedImages()
        self.FixedImages.add()
    
    def show_fixed(self,roi=None,frame=0):
        """ roi and frame are integers.  
        Use Pylab to display the desired ROI (specified by index in self.ROIs).  
        Multiple ROIs can be drawn on separated subplots by passing a list 
        of integers to the function.
        
        """
        fig = pylab.figure()
        
        def add_image(img,sp):
            sp.imshow(img,origin='lower')
            pylab.draw()
        
        # if frame is an integer convert to list
        if type(frame) is type(1):
            frame = [frame]
            
        if type(roi) is type([]):
            for j,f in enumerate(frame):
                img = JT.utils.vtk2numpy(self.FixedImages[f].data.GetImage())
                for i,r in enumerate(roi):
                    print "frame: %i, len(roi): %i, roi: %i, subplot: %i%i%i" %\
                          (f, len(roi) ,r, len(roi),len(frame),r*2+1+f)
                    img_roi = img[self.FixedImages[f].ROIs[r][0][0]:\
                                  self.FixedImages[f].ROIs[r][1][0],\
                                  self.FixedImages[f].ROIs[r][0][1]:\
                                  self.FixedImages[f].ROIs[r][1][1]]
                    sp = fig.add_subplot(len(roi),len(frame),r*2+1+f,
                                         frameon=False, xticks=[], yticks=[])
                    add_image(img_roi,sp)
        else:
            print "Whole image being displayed"
            sp = fig.add_subplot(111)
            add_image(img,sp)

class Volumes(DataNode):
    """
    A container for volumes, based on the built in list type.  There are 
    two ways volumes can be specified, one volume for each object that is being 
    tracked, or one volume with several ROI's.  The first method is probably
    more common for surface models, while the second method would probably be
    more useful for image volumes, especially since the volumes can be very 
    large, and it would be undesirable to keep several copied in memory at 
    the same time
    """
    def __init__(self):
        super(DataNode,self).__init__()
#        self.Calibrations = {}
#        self.Calibrations['external'] = JT.ExternalCalibration()
        # A list of body centers for each ROI in the Volume
        #self.BodyCenters = BodyCenters()
        self.BodyCenters = []
        self.displaySettings = {}
        #self.DisplaySettings.append(None)
        self.data = JT.Volume()
        # self._activeBody is a counter to keep track of the currently 
        # active body for registration
        self.activeBody = 0
        # The number of ROIs contained within the Volume.  The actual ROIs
        # are stored in the Fixed Images branch
        self.nROIs = 1
        # The resultst for 
        # A place to store poses for each body before they are optimized
        # Perhaps the output of user interactor
        self.tempPoses = []
        # A place to store the results of registration for each body
        self.optimizedPoses = []
        
    
class ROIs(DataNode):
    """
    ROI is specified by giving the top left and bottom right image 
    coordinates of the region of interest of the fixed image.
    
    Need to make sure that ROI's and poses are kept in the same index 
    position.
    """
    def __init__(self):
        super(DataNode,self).__init__()
        self.Pose = Pose()
        self.data = None

class Poses(DataNode):
    """
    A container for holding pose data for each volume that is being tracked.
    There should be one pose for each volume/roi that is being tracked.
    The default pose will be set to (0,0,0,0,0,0).
    """
    def __init__(self):
        super(DataNode,self).__init__()
        self.data = Pose()
    
class Calibrations(DataNode):
    def __init__(self):
        super(DataNode,self).__init__()
        self.data = JT.ExternalCalibration()
        self.type = 'external'
        
class FixedImages(DataNode):
    """
    FixedImages is a container for holding fixed images (vtkImageData)
    """
    def __init__(self):
        super(DataNode,self).__init__()
#        self.Calibrations = Calibrations()
#        self.Calibrations.add()
#        self.ROIs = ROIs()
#        self.ROIs.add()
        self.ROIs = ROIs()
        self.data = JT.FixedImage()
        self.calibration_type = 'external'
        self.calibration = JT.ExternalCalibration()
    
    def show(self,roi=None):
        """ roi is an integer.  Use Pylab to display the desired ROI 
        (specified by index in self.ROIs).  Multiple ROIs can be drawn 
        on separated subplots by passing a list of integers the function.
        """
        img = JT.utils.vtk2numpy(self.data.GetImage())
        fig = pylab.figure()
        
        def add_image(img,sp):
            sp.imshow(img,origin='lower')
            pylab.draw()
            
        # if roi is an integer, convert it to a list
        if type(roi) is type(1):
            roi = [roi]
        
        if type(roi) is type([]):
            for i,r in enumerate(roi):
                print "len(roi): ",len(roi), ", roi ", r, ": ", self.ROIs[r]
                img_roi = img[self.ROIs[r][0][0]:self.ROIs[r][1][0],\
                          self.ROIs[r][0][1]:self.ROIs[r][1][1]]
                sp = fig.add_subplot(len(roi),1,i+1,
                                     frameon=False, xticks=[], yticks=[])
                add_image(img_roi,sp)
        else:
            print "Whole image being displayed"
            sp = fig.add_subplot(111)
            add_image(img,sp)

    

class DisplaySettings(DataNode):
    def __init__(self):
        super(DataNode,self).__init__()
        self.data = {}

class BodyCenters(DataNode):
    """Center of body to be used in transform.  This is the point at which
    the volume is rotated about during registration.  It is a 3D position in
    Volume coordinates, relative to the center of the volume.
    """
    def __init__(self):
        super(DataNode,self).__init__()
        self.data = []

#===============================================================================
# Specific Implementations of individual data objects 
#===============================================================================
      
class Pose(JT.vtk.vtkTransform):
    """
    This class will store the pose information for models or volumes that are
    being registered.

    It is subclassed from vtkTransfrom, which describes linear transformations
    as a 4x4 matrix.  Internally, transformations are stored as a 4x4
    homogenous transformation matrix.  Methods are provided for manipulating
    rotation and translations.
    """

    def __init__(self,pose=(0,0,0,0,0,0)):
        """
        pose: a 6 element vector (Position[3], Orientation[3])
        Initialize the pose class.  It will be set to PostMultiply mode.  
        
        """

        self.PostMultiply()
        #This is a bool flag that specifies whether this pose has been obtained
        #maunally or through optimization
        self._isOptimized = False
        self.SetPose(pose)


    def SetPose(self,pose):
        """
        Set the position and orientation of the transform.  Pose is a vector
        made up of (position[3],orientation[3])
        """
        try:
            assert len(pose)==6
        except AssertionError:
            print "Expected 'pose' to contain 6 elements, got %i" % len(pose)
            raise

        self._Position = pose[:3]
        self._Orientation = pose[3:]

        #First reset transformation to identity (otherwise rotations are concatenated)
        self.Identity()

        matrix = self.GetMatrix()
        #Set Position
        for i in range(3):
            matrix.SetElement(i,3,pose[i])

        #Set Rotations
        self.RotateX(pose[3])
        self.RotateY(pose[4])
        self.RotateZ(pose[5])
    
    def GetPose(self):
        """
        Return the position and orientation of the volume as a 6 element
        list of floats. Pose is a vector made up of (position[3],orientation[3])
        """
        position = self.GetPosition()
        orientation = self.GetOrientation()
        pose = [position[0],position[1],position[2],
                orientation[0],orientation[1].orientation[2]]
        return pose
        
        
    def GetPythonArray(self):
        """
        A convienience function for returning a 4x4 matrix as a python numpy array
        """
        pyMatrix = zeros((4,4))
        matrix=self.GetMatrix()
        for i in range(4):
            for j in range(4):
                pyMatrix[i][j]=matrix.GetElement(i,j)

        return pyMatrix

    def Print(self):
        """
        To replicate functionality of JointTrack, it looks like this method
        is used to print the pose data to a text file.
        """
        translation = "%.4f,%.4f,%.4f" % (self.GetPosition())
        rotation = "%.4f,%.4f,%.4f" % (self.GetOrientation())
        return translation + "," + rotation

    def Save(self,file):
        """
        Save (cPickle.dump) data to file object.
        """

    def Load(self):
        """
        Load (cPickle.load) data from file
        """

    def PrintTableHeaders(self):
        """
        """
           
        
def gold_standard():
    Data = JT.Data.Data()
    fixed_image_vert1_filename = []
    fixed_image_vert1_filename.append(
            'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim042-AP.mhd')
    fixed_image_vert1_filename.append(
            'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd')
    fixed_image_vert1_filename.append( 
            r'C:\Users\bryan\bryan-code\trunk\Images\SimulatedFixedImage.mhd')
    fixed_image_vert2_filename = []
    fixed_image_vert2_filename.append(
            'C:/Users/bryan/bryan-code/2D3D/vert2/fluoro/ushortim075-AP.mhd')
    fixed_image_vert2_filename.append(
            'C:/Users/bryan/bryan-code/2D3D/vert2/fluoro/ushortim040-LAT.mhd')
    
    # fixed_image_filename structure: fixed_image_filename[vert][projection]
    fixed_image_filename = [fixed_image_vert1_filename,
                            fixed_image_vert2_filename]
    
    cal_file_vert1 = []
    cal_file_vert1.append(
            'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim042-AP.cal')
    cal_file_vert1.append(
            'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.cal')
    
    cal_file_vert2 = []
    cal_file_vert2.append(
            'C:/Users/bryan/bryan-code/2D3D/vert2/fluoro/ushortim075-AP.cal')
    cal_file_vert2.append(
            'C:/Users/bryan/bryan-code/2D3D/vert2/fluoro/ushortim040-LAT.cal')
    
    # cal_file structure: cal_file[vert][proj]
    cal_file = [cal_file_vert1,cal_file_vert2]
    
    volume_filename = []
    volume_filename.append(
            'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mhd')
    volume_filename.append(
            'C:/Users/bryan/bryan-code/2D3D/vert2/CT/ucharCTv2r4.mhd')
    
    sta_file = []
    sta_file.append('C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta')
    sta_file.append('C:/Users/bryan/bryan-code/2D3D/vert2/CT/ucharCTv2r4.sta')
    
    roi_vert1_ap = []
    roi_vert1_ap.append([[268,166],[344,313]])
    roi_vert1_ap.append([[209,175],[283,341]])
    roi_vert1_ap.append([[155,183],[223,331]])
    
    roi_vert1_lat = []
    roi_vert1_lat.append([[273,204],[343,371]])
    roi_vert1_lat.append([[210,187],[277,370]])
    roi_vert1_lat.append([[160,190],[223,347]])
    
    roi_sim_lat = []
    roi_sim_lat.append([[273,204],[343,371]])
    roi_sim_lat.append([[210,187],[277,370]])
    roi_sim_lat.append([[160,190],[223,347]])
    
    roi_vert2_ap = []
    roi_vert2_ap.append([[332,205],[403,369]]) #vert 2, body 1
    roi_vert2_ap.append([[275,210],[345,387]]) #vert 2, body 2
    roi_vert2_ap.append([[213,227],[291,371]]) #vert 2, body 3
    roi_vert2_ap.append([[154,219],[221,380]]) #vert 2, body 4
    roi_vert2_ap.append([[86,220],[156,366]]) #vert 2, body 5
    
    roi_vert2_lat = []
    roi_vert2_lat.append([[337,166],[394,299]]) #vert 2, body 1
    roi_vert2_lat.append([[273,156],[341,282]]) #vert 2, body 2
    roi_vert2_lat.append([[217,146],[284,286]]) #vert 2, body 3
    roi_vert2_lat.append([[152,129],[229,285]]) #vert 2, body 4
    roi_vert2_lat.append([[89,107],[170,282]]) #vert 2, body 5
    
    # roi structure: roi[vert][projection][body]
    roi_vert1 = [roi_vert1_ap,roi_vert1_lat,roi_sim_lat]
    roi_vert2 = [roi_vert2_ap,roi_vert2_lat]
    roi = [roi_vert1,roi_vert2]
    
    centers_vert1 = []
    centers_vert1.append([3.906, -24.738, -12.586])
    centers_vert1.append([3.906, 2.17, -7.378])
    centers_vert1.append([7.378, 27.342, -6.51])
    
    centers_vert2 = []
    centers_vert2.append([8.07674, -41.94694, -14.85078])
    centers_vert2.append([3.9081, -21.10374, -17.45618])
    centers_vert2.append([-0.26054, 0.78162, -21.10374])
    centers_vert2.append([-5.47134, 22.66698, -25.79346])
    centers_vert2.append([-10.16106, 45.07342, -27.87778])
    
    # centers structure: centers[vert][body]
    centers = [centers_vert1,centers_vert2]
    
    projection_range = range(len(fixed_image_filename))
    vert_range = range(len(volume_filename))
    center_range = []
    for vert in vert_range:
        center_range.append(range(len(centers[vert])))    
    
    
    # Frame 1 is vert 1, Frame 2 is vert 2  
    # Load Data elements into Data structure
    for vert in vert_range:
        Data.Frames.add()
        Data.Frames[vert].Volumes[0].data.\
            SetVolumeFilename(volume_filename[vert])
            
        for proj in projection_range:
            Data.Frames[vert].FixedImages.add()
            for body in center_range[vert]:
                #print "vert: %i, proj: %i, body: %i" % (vert, proj, body)
                #print "centers: ", centers[vert][body]
                Data.Frames[vert].Volumes[0].BodyCenters.append(centers[vert][body])
                
                Data.Frames[vert].FixedImages[proj].data.\
                    SetFileName(fixed_image_filename[vert][proj])
                
                Data.Frames[vert].FixedImages[proj].calibration.\
                    LoadCalFile(cal_file[vert][proj])
                
                Data.Frames[vert].FixedImages[proj].calibration.\
                    LoadStaFile(sta_file[vert])
                
                Data.Frames[vert].Volumes[0].data.SetOriginalTransform(
                        Data.Frames[vert].FixedImages[proj].calibration.
                        volumeTransform)
                  
                Data.Frames[vert].FixedImages[proj].ROIs.add()
                Data.Frames[vert].FixedImages[proj].ROIs[body].data =\
                        roi[vert][proj][body]
        
    return Data

def simulated_data():
    Data = JT.Data.Data()
    fixed_image_filename = \
            r'C:\Users\bryan\bryan-code\trunk\Tests\SimulatedFixedImage.mhd'
    
    cal_file = \
            'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.cal'
    
    volume_filename = \
            'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mhd'
    
    sta_file = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
    
    roi = []
    roi.append([[273,204],[343,371]])
    roi.append([[210,187],[277,370]])
    roi.append([[160,190],[223,347]])
    
    centers = []
    centers.append([3.906, -24.738, -12.586])
    centers.append([3.906, 2.17, -7.378])
    centers.append([7.378, 27.342, -6.51])
    
    body_range = range(len(centers))
 
    # Load Data elements into Data structure
    Data.Frames.add()
    Data.Frames[0].Volumes[0].data.\
        SetVolumeFilename(volume_filename)
        
    Data.Frames[0].FixedImages.add()
    for body in body_range:
        #print "vert: %i, proj: %i, body: %i" % (vert, proj, body)
        #print "centers: ", centers[vert][body]
        Data.Frames[0].Volumes[0].BodyCenters.append(centers[body])
        Data.Frames[0].FixedImages[0].data.SetFileName(fixed_image_filename)
        Data.Frames[0].FixedImages[0].calibration.LoadCalFile(cal_file)
        Data.Frames[0].FixedImages[0].calibration.LoadStaFile(sta_file)
        Data.Frames[0].Volumes[0].data.SetOriginalTransform(
                Data.Frames[0].FixedImages[0].calibration.volumeTransform)
        Data.Frames[0].FixedImages[0].ROIs.add()
        Data.Frames[0].FixedImages[0].ROIs[body].data = roi[body]
    
    return Data
    
if __name__ == '__main__':
    gold_standard()       
